package com.sunzm.spark.streaming

import org.apache.spark.SparkConf
import org.apache.spark.streaming.dstream.{DStream, MapWithStateDStream}
import org.apache.spark.streaming.{Durations, State, StateSpec, StreamingContext}

object MapWithStateDemo {
  private val checkpointDirectory = "data/spark/streaming/ck/20210612002/"

  // Function to create and setup a new StreamingContext
  def functionToCreateContext(): StreamingContext = {
    val conf = new SparkConf()
      .setAppName(s"${this.getClass.getSimpleName.stripSuffix("$")}")
      .setMaster("local[*]")

    val ssc = new StreamingContext(conf, Durations.seconds(5))

    //设置checkpoint目录
    ssc.checkpoint(checkpointDirectory)

    processFunc(ssc)

    ssc
  }

  def main(args: Array[String]): Unit = {
    //从checkpoint目录中恢复StreamingContext，如果不存在，就调用 functionToCreateContext 方法 创建一个新的StreamingContext
    val ssc = StreamingContext.getOrCreate(checkpointDirectory, functionToCreateContext _)

    ssc.start()
    ssc.awaitTermination()
  }

  def processFunc(ssc: StreamingContext) = {
    val lines = ssc.socketTextStream("82.156.210.70", 9999)

    val pairsRDD: DStream[(String, Int)] = lines.flatMap(_.split(","))
      .map((_, 1))

    val spec: StateSpec[String, Int, Int, String] = StateSpec.function(mappingFunc)

    val mapWithStateDStream: MapWithStateDStream[String, Int, Int, String] = pairsRDD.mapWithState(spec)

    //mapWithStateDStream就是处理后的结果数据
    mapWithStateDStream.foreachRDD(rdd => {
      if(!rdd.isEmpty()){
        rdd.foreach(line => {
          println(s"结果数据: ${line}")
        })
      }else{
        println(s"结果数据为空!")
      }
    })

    //stateSnapshots返回的是包含key的状态数据
    val stateDStream: DStream[(String, Int)] = mapWithStateDStream.stateSnapshots()

    stateDStream.foreachRDD(rdd => {
      if(!rdd.isEmpty()){
        rdd.foreach(line => {
          println(s"状态数据: ${line}")
        })
      }else{
        println(s"状态数据为空!")
      }
    })
  }

  /**
   * mapWithState 只有K-V类型的DStream才能调用
   *
   * StateSpec.function需要传递的函数的定义如下：
   * mappingFunction: (KeyType, Option[ValueType], State[StateType]) => MappedType
   *
   * KeyType是key的类型
   * Option[ValueType]: ValueType是value的类型
   * State[StateType]: StateType 是状态的数据类型
   * MappedType 是返回值的数据类型
   *
   * 下面这个例子中
   * word 是key
   * one 是一条数据的value
   * state 是状态数据
   */
  val mappingFunc = (word: String, one: Option[Int], state: State[Int]) => {
    val sum = one.getOrElse(0) + state.getOption.getOrElse(0)
    val output = s"单词:${word}, 数量: ${sum}"
    state.update(sum)
    output
  }
}
